import M2Crypto
from M2Crypto import *
import base64
import os
import sys
import datetime
import tempfile
from binascii import hexlify, unhexlify
import random

#####################################################################
####    UTILITIES
#####################################################################

def plus_twentyyears(dt):
    """helper function to return a datetime with some 20 years added to it
    """
    return datetime.datetime(dt.year+20, dt.month, dt.day,
            dt.hour, dt.minute, dt.second, dt.microsecond)

def plus_oneyear(dt):
    """helper function to return a datetime with some 20 years added to it
    """
    return datetime.datetime(dt.year+1, dt.month, dt.day,
            dt.hour, dt.minute, dt.second, dt.microsecond)


def generate_nonce(message):

    nonce = hexlify( EVP.hmac(message, str(random.random()), algo='sha256'))
    return nonce



def datetime_to_timestamp(dt):
    """this function takes a datetime.datetime instance as parameter
    and returns an integer representing the number for seconds
    since the 1st of january 1970 (epoch).

    WARNING: OpenSSL considers the epoch to be 1/1/1970 and thus
    this function does the same. Do not try to use this function
    with datetimes that are older than this date.
    """
    delta = dt - epoch

    return delta.days * 24 * 3600 + delta.seconds

def timestamp_to_datetime(ts):
    """this function takes a timestamp expressed with 1/1/1970 as the
    epoch and returns a python datetime.datetime instance
    """
    delta = datetime.timedelta(0, ts)

    return epoch + delta


def get_secure_filename():
    """Creates a tempfile in the most secure manner possible,
    make sure is it closed and return the filename for
    easy usage.
    """

    file_handle, filename = tempfile.mkstemp()
    tmpfile = os.fdopen(file_handle, "r")
    tmpfile.close()
    return filename


#--------- Function for keys ------------------------------------

PASSWORD = "nopass"
def call(arg):
    print "Using Pass for keys : ", PASSWORD
    return PASSWORD


def askForPass ( ) :
    passWord = raw_input("Give password for keys :\n")
    global PASSWORD
    PASSWORD = passWord
    return passWord

#----------------------------------------------------------------

#####################################################################

now = datetime.datetime.now

PADDING_TYPE = M2Crypto.RSA.pkcs1_padding

def create_cert_request( pubkey, privkey, cert_name, digest="md5"):
    """Create a certificate request.

    @param pubkey: the EVP.PKey public key to associate with the request
    @type pubkey: EVP.PKey instance

    @param privkey: the EVP.PKey public key to associate with the request
    @type privkey: EVP.PKey instance

    @param digest: Digest method to use for signing, default is md5
    @type digest: string

    @param: cert_name - The name of the subject of the request, possible
    arguments are:
    C     - Country name
    ST    - State or province name
    L     - Locality name
    O     - Organization name
    OU    - Organizational unit name
    CN    - Common name
    emailAddress - E-mail address

    @returns: The certificate request in an X509.X509Req object
    """
    req = X509.Request()


    req.set_subject_name(cert_name)
    req.set_pubkey(pubkey)
    req.sign(privkey, digest)
    
    s = req.get_subject()

    

    return req

def create_certificate(req, (issuerCert, issuerKey), serial,
        (notbefore, notafter), digest="sha1"):
    """
    Generate a certificate given a certificate request.

    Arguments: req        - Certificate reqeust to use
               issuerCert - The certificate of the issuer
               issuerKey  - The private key of the issuer
               serial     - Serial number for the certificate
               notbefore  - a datetime.datetime instance when the certificate
                            starts being valid
               notafter   - a datetime.datetime instance when the certificate
                            stops being valid
               digest     - Digest method to use for signing, default is md5
    Returns:   The signed certificate in an X509 object

    WARNING: the datetime instances passed for notbefore and notafter should
    be expressed in UTC.
    """

    not_before = ASN1.ASN1_UTCTIME()
    not_before.set_time(datetime_to_timestamp(notbefore))

    not_after = ASN1.ASN1_UTCTIME()
    not_after.set_time(datetime_to_timestamp(notafter))

    cert = X509.X509()
    cert.set_serial_number(serial)
    
    cert.set_not_before(not_before)
    cert.set_not_after(not_after)
    
    cert.set_issuer(issuerCert.get_subject())

    cert.set_subject(req.get_subject())
    cert.set_pubkey(req.get_pubkey())
    cert.sign(issuerKey, digest) 

    
    return cert

def as_evp_pkey(rsa_key):
    """Transform an rsa key into an evp_pkey instance
    """
    evp_pkey = EVP.PKey()
    success = evp_pkey.assign_rsa(rsa_key, capture=0)
    # TODO: Test the success
    return evp_pkey



def create_rsa_keyring(key_size=2048, exponent=65537, callback=call):
    """ Return a new RSA keyring as tuple (<public key>, <private key>)
    @param passphrase_callback: A Python callable object that is invoked
    during key generation; its usual purpose is to provide visual
    feedback.
    You must load a randpool file before calling this method with
    Rand.load_file('randpool', -1).
    After calling the method, you must save the randpool file with
    Rand.save_file('randpool')
    @type passphrase_callback: Callback

    @param key_size: Key length, in bits
    @type key_size: Integer

    @param exponent: The RSA public exponent
    @type exponent: Integer
    """
    keyring = RSA.gen_key(key_size, exponent)

    privkey_filename = get_secure_filename()
    pubkey_filename = get_secure_filename()

    #keyring.save_key(privkey_filename, callback=util.passphrase_callback)
    keyring.save_key(privkey_filename, callback=call)
    keyring.save_pub_key(pubkey_filename)

    #privkey = RSA.load_key(privkey_filename, callback=util.passphrase_callback)
    privkey = RSA.load_key(privkey_filename, callback)
    pubkey = RSA.load_pub_key(pubkey_filename)

    # This is evp keystore
    evp_pubkey = EVP.load_key(privkey_filename, callback)

    os.unlink(privkey_filename)
    os.unlink(pubkey_filename)

    return (pubkey, privkey, evp_pubkey)


def create_rsa_keyring_save(file, key_size=2048, exponent=65537, callback=call):
    """ Return a new RSA keyring as tuple (<public key>, <private key>)
    @param passphrase_callback: A Python callable object that is invoked
    during key generation; its usual purpose is to provide visual
    feedback.
    You must load a randpool file before calling this method with
    Rand.load_file('randpool', -1).
    After calling the method, you must save the randpool file with
    Rand.save_file('randpool')
    @type passphrase_callback: Callback

    @param key_size: Key length, in bits
    @type key_size: Integer

    @param exponent: The RSA public exponent
    @type exponent: Integer
    """
    keyring = RSA.gen_key(key_size, exponent)

    privkey_filename = get_secure_filename()
    pubkey_filename = get_secure_filename()

    #keyring.save_key(privkey_filename, callback=util.passphrase_callback)
    keyring.save_key(privkey_filename, callback=call)
    keyring.save_key(file, callback=call)
    keyring.save_pub_key(pubkey_filename)

    #privkey = RSA.load_key(privkey_filename, callback=util.passphrase_callback)
    privkey = RSA.load_key(privkey_filename, callback)
    pubkey = RSA.load_pub_key(pubkey_filename)

    # This is evp keystore
    evp_pubkey = EVP.load_key(privkey_filename, callback)

    os.unlink(privkey_filename)
    os.unlink(pubkey_filename)

    return (pubkey, privkey, evp_pubkey)


def load_rsa_keyring(file, key_size=2048, exponent=65537, callback=call):
    """ Return a new RSA keyring as tuple (<public key>, <private key>)
    @param passphrase_callback: A Python callable object that is invoked
    during key generation; its usual purpose is to provide visual
    feedback.
    You must load a randpool file before calling this method with
    Rand.load_file('randpool', -1).
    After calling the method, you must save the randpool file with
    Rand.save_file('randpool')
    @type passphrase_callback: Callback

    @param key_size: Key length, in bits
    @type key_size: Integer

    @param exponent: The RSA public exponent
    @type exponent: Integer
    """
    #privkey = RSA.load_key(privkey_filename, callback=util.passphrase_callback)
    privkey = RSA.load_key(file, callback)
#    pubkey = RSA.load_pub_key(file)
    pubkey = privkey;

    # This is evp keystore
    evp_pubkey = EVP.load_key(file, callback)


    return (pubkey, privkey, evp_pubkey)

def create_certificate_authority(passphrase_callback=call, serial=0, cert_name=None,
        notbefore=now(), notafter=plus_twentyyears(now()), digest="sha1", file="tmp"):
    """Function to create a certificate authority.

    @param passphrase_callback: a callback that will be called without
    argument and that must return the passphrase you want to use
    on your CA key
    @type passphrase_callback: callable

    @param serial: Serial number for the certificate, default is 0
    @type serial: Integer

    @param notbefore: datetime.datetime instance when
    the certificate starts being valid, default is now
    @type notbefore: datetime.datetime instance

    @param notafter: datetime.datetime instance when the certificate
    stops being valid, default is 20 years more than now
    @type notafter: datetime.datetime instance

    @param: **name - The name of the subject of the request, possible
    arguments are:
    C     - Country name
    ST    - State or province name
    L     - Locality name
    O     - Organization name
    OU    - Organizational unit name
    CN    - Common name
    emailAddress - E-mail address

    """
    assert isinstance(notbefore, datetime.datetime)
    assert isinstance(notafter, datetime.datetime)
    pubkey, privkey, ekey = create_rsa_keyring_save(file+".key")

    ca_priv_evp = as_evp_pkey(privkey)
    ca_pub_evp = as_evp_pkey(pubkey)
    ca_cert_req = create_cert_request(ca_pub_evp, ca_priv_evp, cert_name, digest)

    ca_cert = create_certificate(
            ca_cert_req,
            (ca_cert_req, ca_priv_evp),
            serial,
            (notbefore, notafter)
            )

    return ca_cert, pubkey, privkey, ekey


def load_certificate_authority(passphrase_callback=call, serial=0, cert_name=None,
        notbefore=now(), notafter=plus_twentyyears(now()), digest="sha1", file="tmp"):
    """Function to load a certificate authority.

    @param passphrase_callback: a callback that will be called without
    argument and that must return the passphrase you want to use
    on your CA key
    @type passphrase_callback: callable

    @param serial: Serial number for the certificate, default is 0
    @type serial: Integer

    @param notbefore: datetime.datetime instance when
    the certificate starts being valid, default is now
    @type notbefore: datetime.datetime instance

    @param notafter: datetime.datetime instance when the certificate
    stops being valid, default is 20 years more than now
    @type notafter: datetime.datetime instance

    @param: **name - The name of the subject of the request, possible
    arguments are:
    C     - Country name
    ST    - State or province name
    L     - Locality name
    O     - Organization name
    OU    - Organizational unit name
    CN    - Common name
    emailAddress - E-mail address

    """
    assert isinstance(notbefore, datetime.datetime)
    assert isinstance(notafter, datetime.datetime)
    pubkey, privkey, ekey = load_rsa_keyring(file+".key")


    ca_cert = X509.load_cert(file+".pem")

    return ca_cert, pubkey, privkey, ekey


def load_certificate_client(file, file_ca,  passphrase_callback=call, serial=0, cert_name=None,
        notbefore=now(), notafter=plus_twentyyears(now()), digest="sha1" ):
    """Function to load a certificate authority.

    @param passphrase_callback: a callback that will be called without
    argument and that must return the passphrase you want to use
    on your CA key
    @type passphrase_callback: callable

    @param serial: Serial number for the certificate, default is 0
    @type serial: Integer

    @param notbefore: datetime.datetime instance when
    the certificate starts being valid, default is now
    @type notbefore: datetime.datetime instance

    @param notafter: datetime.datetime instance when the certificate
    stops being valid, default is 20 years more than now
    @type notafter: datetime.datetime instance



    """
    assert isinstance(notbefore, datetime.datetime)
    assert isinstance(notafter, datetime.datetime)
    pubkey, privkey, ekey = load_rsa_keyring(file+".key")

    cert = X509.load_cert(file+".pem")
    ca_cert = X509.load_cert(file_ca+".pem")


    return cert, privkey, pubkey, ekey, ca_cert;


    return ca_cert, pubkey, privkey, ekey
def create_x509_certificate(ca_cert, ca_privkey, cert_name,
        cert_passphrase_callback, ca_passphrase_callback, serial=0, notbefore=now(),
        notafter=plus_twentyyears(now())):
    """Create and return a new x509 certificate using a Certificate Authority

    @param ca_cert: The Certificate Authority to use to create
    the new X509 certificate
    @type ca_cert: M2Crypto.X509.X509 instance

    @param ca_privkey: The Certificate Authority private key, its the
    private key which was used for sign the CA
    @type ca_privkey: M2Crypto.RSA.RSA instance

    @param cert_passphrase_callback: A Python callable object that is invoked
    during key generation; its usual purpose is to provide visual
    feedback.
    @type cert_passphrase_callback: Callback

    @param ca_passphrase_callback:  A Python callable object that is invoked
    during the new certificate signature; its usual purpose is to provide visual
    feedback.

    @param serial: Serial number for the certificate, default is 0
    @type serial: Integer

    @param notbefore: datetime.datetime instance when the certificate
    starts being valid, default is 0
    @type notbefore: datetime.datetime instance

    @param notafter: datetime.datetime instance when the certificate
    stops being valid, default is 20 years
    @type notafter: datetime.datetime instance

    @param: **name - The name of the subject of the request, possible
    arguments are:
    C     - Country name
    ST    - State or province name
    L     - Locality name
    O     - Organization name
    OU    - Organizational unit name
    CN    - Common name
    emailAddress - E-mail address
    """
   
    pubkey, privkey, evp_key = create_rsa_keyring()

    priv_evp = as_evp_pkey(privkey)
    pub_evp = as_evp_pkey(pubkey)
    ca_privkey_evp = as_evp_pkey(ca_privkey)
    cert_req = create_cert_request(pub_evp, priv_evp, cert_name)
    cert = create_certificate(
            cert_req,
            (ca_cert, ca_privkey_evp),
            serial,
            (notbefore, notafter)
            )

    return cert, pubkey, privkey

def verify_x509_certificate(ca_cert, client_cert):
    """Verify if a certificate was signed by a certain CA.
    Return 1, if the CA certificate is ok otherwise 0

    @param ca_cert: The CA certificate to use to verify the client certificate
    @type ca_cert: M2Crypto.X509.X509 instance

    @param client_cert: The client certificate to verify
    @type client_cert: M2Crypto.X509.X509 instance

    @result: Return an integer 1 or 0
    """
    success = client_cert.verify(ca_cert.get_pubkey())

    return success

def encrypt(pubkey, clear_data):
    """Encrypt <clear_data> string with the public key from
    the certificate
    You must load a randpool file before calling this method with
    Rand.load_file('randpool', -1).
    After calling the method, you must save the randpool file
    with Rand.save_file('randpool')
    
    @param pubkey: The public key to use to encrypt the data,
    if you have an X509 certificate,
    you use the public key from certificate.get_pubkey() method
    to retrieve public key
    @type pubkey: M2Crypto.EVP.PKey instance

    @param clear_data: Clear data to encrypt
    @type clear_data: String

    @return: the encrypted data encoded with base64
    """
    rsapub = RSA.RSA_pub(pubkey.get_rsa().rsa)
    encrypted_data = rsapub.public_encrypt(clear_data, PADDING_TYPE)
    return base64.b64encode(encrypted_data)

def decrypt(privkey, encrypted_data):
    """Decrypt <encrypt_data> with the private key
    Encrypted data must be encode with base64
    
    @param privkey: The private RSA key to use for decryption
    @type privkey: M2Crypto.X509.X509 instance

    @param encrypt_data: Encrypted data (String) to decrypt,
    the encrypted data must be encoded with base64
    @type encrypt_data: String

    @return: the decrypted (clear) data
    """
    k = base64.b64decode(encrypted_data)
    clear_data = privkey.private_decrypt(k, PADDING_TYPE)
    return clear_data

def sign(privkey, data):
    """Sign the data and return the generated signature encoding with base64
    You must load a randpool file before calling this method with
    Rand.load_file('randpool', -1).
    After calling the method, you must save the randpool file with
    Rand.save_file('randpool')
    
    @param privkey: The private RSA key to use to sign the data
    @type privkey: M2Crypto.RSA.RSA instance

    @param data: The data string to sign
    @type data: String
    """
    priv_evp = as_evp_pkey(privkey)
    priv_evp.sign_init()
    priv_evp.update(data)
    sig = priv_evp.final()
    result = base64.b64encode(sig)
    return result
    
def sign_file(privkey, file_path):
    """Sign the data and return the generated signature encoding with base64
    
    @param privkey: The private RSA key to use to sign the data
    @type privkey: M2Crypto.RSA.RSA instance

    @param file_path: The path of file that contain data to sign
    @type file_path: String
    """
    priv_evp = as_evp_pkey(privkey)
    priv_evp.sign_init()
    f = open(file_path, 'rb')
    while True:
        buf = f.read(4096)
        if buf == '':
            break
        priv_evp.update(buf)

    f.close()
    sig = priv_evp.final()
    result = base64.b64encode(sig)

    return result

def verify(pubkey, data, signature):
    """Verify the data with the signature.
    signature must be encoded with base64

    @param pubkey: The public key to use to verify the signature with the data,
    if you have an X509 certificate, you use the public key from
    certificate.get_pubkey() method to retrieve public key
    @type pubkey: M2Crypto.EVP.PKey instance

    @param data: The data to verify with the signature
    @type file_path: String

    @param signature: The signature, the signature must be encoded in base64
    @type signature: String

    @result: Return 1 otherwise the method return an exception
    """
    M_sig = base64.b64decode(signature)
    pubkey.verify_init()
    pubkey.verify_update(data)
    result = pubkey.verify_final(M_sig)
    return result

def verify_file(pubkey, file_path, signature):
    """Verify the data with the signature
    signature must be encoded with base64

    @param pubkey: The public key to use to verify the signature with the data,
    if you have an X509 certificate, you use the public key from
    certificate.get_pubkey() method to retrieve public key
    @type pubkey: M2Crypto.EVP.PKey instance

    @param file_path: The path of the file that contain data to verify
    with the signature
    @type file_path: String

    @param signature: The signature, the signature must be encoded in base64
    @type signature: String

    @result: Return 1 otherwise the method return an exception
    """
    M_sig = base64.b64decode(signature)

    pubkey.verify_init()
    f = open(file_path, 'rb')
    while True:
        buf = f.read(4096)
        if buf == '':
            break
        pubkey.verify_update(buf)

    f.close()
    result = pubkey.verify_final(M_sig)

    return result

epoch = datetime.datetime(1970, 1, 1, 0, 0, 0)

